///////////////////////////////////////////////////////////////////////
// Project: XNA Quake3 Lib - BSP 
// Author: Craig Sniffen
// Copyright (c) 2008-2009 All rights reserved
///////////////////////////////////////////////////////////////////////

using System;
using System.IO;
using System.Text;
using System.Collections.Generic;
using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Graphics;
using Microsoft.Xna.Framework.Content.Pipeline;
using Microsoft.Xna.Framework.Content.Pipeline.Graphics;
using XNAQ3Lib.Q3BSP;

// TODO: replace this with the type you want to import.
using TImport = System.Collections.Generic.Dictionary<string, XNAQ3Lib.Q3BSP.Q3BSPMaterial>;

namespace Q3BSPContentPipelineExtension
{
    /// <summary>
    /// This class will be instantiated by the XNA Framework Content Pipeline
    /// to import a file from disk into the specified type, TImport.
    /// 
    /// This should be part of a Content Pipeline Extension Library project.
    /// 
    /// TODO: change the ContentImporter attribute to specify the correct file
    /// extension, display name, and default processor for this importer.
    /// </summary>
    [ContentImporter(".shader", DisplayName = "XNA Q3 Library Shader Importer")]
    public class ShaderContentImporter : ContentImporter<TImport>
    {
        public override TImport Import(string filename, ContentImporterContext context)
        {
            //System.Diagnostics.Debugger.Launch();

            TImport shaderDictionary = new TImport();
            StreamReader sr = new StreamReader(filename);
            StreamWriter sw = new StreamWriter(filename.Substring(0, filename.LastIndexOf('.')) + ".compiled.txt");
            
            string shaderName = "";
            Q3BSPMaterial shader = new Q3BSPMaterial();
            Q3BSPMaterialStage currentStage = new Q3BSPMaterialStage();
            bool insideShaderBlock = false;
            bool insideStageBlock = false;
            shader.stages = new List<Q3BSPMaterialStage>();

            while(!sr.EndOfStream)
            {
                string line = sr.ReadLine();
                int comment;

                // Remove comments and whitespace
                if ((comment = line.LastIndexOf("//")) > -1)
                    line = line.Substring(0, comment);
                line = line.Trim();
                
                // Line is empty don't process it
                if (line == "")
                    continue;

                // Line is important only when the BSP is compiled, don't process it
                if (line.Length > 3 && line.Substring(0, 4) == "qer_")
                    continue;
                if (line.Length > 4 && line.Substring(0, 5) == "q3map")
                    continue;
                if (line.Length > 10 && line.Substring(0, 11) == "surfaceparm")
                    continue;

                #region Read lines outside any blocks
                if (!insideShaderBlock)
                {
                    if (line.Substring(0, 1) == "{")
                    {
                        if (shaderName == "")
                            throw new Exception("Error loading .shader: there is a block without a shader name.");
                        
                        shader = new Q3BSPMaterial();
                        shader.shaderName = shaderName;

                        insideShaderBlock = true;
                    }

                    else
                    {
                        if (shaderName != "")
                            throw new Exception("Error loading .shader: Shader " + shaderName + " has no information describing it.");
                        
                        shaderName = line;
                    }
                }
                #endregion

                #region Read lines inside the shader block but outside the stage blocks
                else if (insideShaderBlock && !insideStageBlock)
                {
                    if (line.Substring(0, 1) == "{")
                    {
                        insideStageBlock = true;
                        currentStage = new Q3BSPMaterialStage();

                    }

                    else if (line.Length > 8 && line.Substring(0, 8) == "skyparms")
                    {

                        ProcessSkyparms(line.Substring(8).Trim(), ref shader);
                    }

                    else if (line.Substring(0, 1) == "}")
                    {
                        if (shader.stages.Count < 1)
                            throw new Exception("Error loading .shader: Shader " + shader.shaderName + " has no stages.");

                        shader.ShaderCode = Effect.CompileEffectFromSource(GenerateEffectFromShader(ref shader, ref sw), null, null, CompilerOptions.None, TargetPlatform.Windows).GetEffectCode();
                        shaderDictionary.Add(shader.shaderName, shader);

                        shader = null;
                        shaderName = "";
                        insideShaderBlock = false;
                    }
                }
                #endregion

                #region Read lines inside stage blocks
                else if (insideStageBlock)
                {
                    if (line.Length > 2 && line.Substring(0, 3) == "map")
                    {
                        string textureFilename = line.Substring(4);

                        // Account for special texures (ie $lightmap)
                        if (textureFilename.Substring(0, 1) == "$")
                        {
                            if(String.Compare(textureFilename, "$lightmap", true) == 0)
                            {
                                currentStage.IsLightmapStage = true;
                            }
                            textureFilename = textureFilename.Substring(1);
                        } 
                        
                        if (textureFilename.LastIndexOf('.') > -1)
                        {
                            textureFilename = textureFilename.Substring(0, textureFilename.LastIndexOf('.'));
                        }

                        currentStage.TextureFilename = textureFilename;
                        
                    }

                    else if (line.Length > 4 && line.Substring(0, 5) == "tcMod")
                    {
                        ProcessTcMod(line.Substring(5).Trim(), ref currentStage);
                    }

                    else if (line.Substring(0, 1) == "}")
                    {
                        shader.stages.Add(currentStage);
                        insideStageBlock = false;
                    }
                }
                #endregion
            }
            
            sw.Dispose();

            return shaderDictionary;
        }



        string GenerateEffectFromShader(ref Q3BSPMaterial shader, ref StreamWriter sw)
        {
            StringBuilder sb = new StringBuilder();
            sw.WriteLine("Shader: " + shader.shaderName);

            int stageNumber = 0;
            int texCoodNumber = 0;

            string externs = "";
            string samplers = "";
            string vertexShaderInput = "";
            string vertexShaderOutput = "";
            string stageInputStructs = "";
            string vertexShaderProgram = "";
            string pixelShaderProgram = "";
            string technique = "";

            #region Generate Genertic External Variables
            sb.AppendLine("// Generic Externs");
            sb.AppendLine("uniform float4x4 worldViewProj;");

            externs += sb.ToString();
            sb.Remove(0, sb.ToString().Length);
            #endregion

            foreach (Q3BSPMaterialStage stage in shader.stages)
            {
                stage.InitializeEffectNames(stageNumber);

                #region Generate External Variables
                sb.AppendLine("// " + stage.stageName + "'s externs");
                sb.AppendLine("uniform texture " + stage.textureName + ";");
                sb.AppendLine("uniform float3x2 " + stage.tcModName + ";");
                externs += sb.ToString();

                sb.Remove(0, sb.ToString().Length);
                #endregion

                #region Generate Samplers
                sb.AppendLine("sampler " + stage.samplerName + " = sampler_state");
                sb.AppendLine("{");
                sb.AppendLine("\tTexture = <" + stage.textureName + ">;");
                sb.AppendLine("\tminfilter = LINEAR;");
                sb.AppendLine("\tmagfilter = LINEAR;");
                sb.AppendLine("\tmipfilter = NONE;");
                sb.AppendLine("\tAddressU = Wrap;");
                sb.AppendLine("\tAddressV = Wrap;");
                sb.AppendLine("};");
                samplers += sb.ToString();

                sb.Remove(0, sb.ToString().Length);
                #endregion

                #region Generate Input Structures
                sb.AppendLine("struct " + stage.inputName);
                sb.AppendLine("{");
                sb.AppendLine("\tfloat2 uv : TEXCOORD" + texCoodNumber + ";");
                sb.AppendLine("};");
                stageInputStructs += sb.ToString();

                sb.Remove(0, sb.ToString().Length);

                texCoodNumber++;
                #endregion

                stageNumber++;
            }

            #region Generate VertexShaderInput
            sb.AppendLine("struct VertexShaderInput");
            sb.AppendLine("{");
            if (shader.IsSky)
            {
                sb.AppendLine("\tfloat4 position : POSITION0;");
            }
            else
            {
                sb.AppendLine("\tfloat4 position : POSITION0;");
                sb.AppendLine("\tfloat3 normal : NORMAL0;");
                sb.AppendLine("\tfloat2 texcoords : TEXCOORD0;");
                sb.AppendLine("\tfloat2 lightmapcoords : TEXCOORD1;");
                sb.AppendLine("\tfloat4 diffuse : COLOR0;");
            }
            sb.AppendLine("};");
            vertexShaderInput = sb.ToString();

            sb.Remove(0, sb.ToString().Length);
            #endregion

            #region Generate VertexShaderOutput
            sb.AppendLine("struct VertexShaderOutput");
            sb.AppendLine("{");
            sb.AppendLine("\tfloat4 position : POSITION0;");
            foreach (Q3BSPMaterialStage stage in shader.stages)
            {
                sb.AppendLine("\t" + stage.inputName + " " + stage.stageName + ";");
            }
            sb.AppendLine("};");
            vertexShaderOutput = sb.ToString();

            sb.Remove(0, sb.ToString().Length);
            #endregion

            #region Generate VertexShaderProgram
            sb.AppendLine("VertexShaderOutput VertexShaderProgram(VertexShaderInput input)");
            sb.AppendLine("{");
            sb.AppendLine("\tVertexShaderOutput output;");
            sb.AppendLine("\toutput.position = mul(input.position, worldViewProj);");

            if (shader.IsSky)
            {
                float zScale = shader.SkyHeight/256.0f;
                float zOffset = 0.707f;

                sb.AppendLine("\tfloat3 S = normalize(float3(input.position.x, input.position.z, input.position.y));");
                sb.AppendLine("\tS.z = " + zScale + " * (S.z + " + zOffset + ");");
                sb.AppendLine("\tS = normalize(S);");
                sb.AppendLine();

                foreach (Q3BSPMaterialStage stage in shader.stages)
                {
                    sb.AppendLine("\toutput." + stage.stageName + ".uv = mul(float3(-S.x, -S.y, 1.0), " + stage.tcModName + ");");
                }
            }
            else
            {
                foreach (Q3BSPMaterialStage stage in shader.stages)
                {
                    if (stage.IsLightmapStage)
                    {
                        sb.AppendLine("\toutput." + stage.stageName + ".uv = mul(float3(input.lightmapcoords, 1.0), " + stage.tcModName + ");");

                    }
                    else
                    {
                        sb.AppendLine("\toutput." + stage.stageName + ".uv = mul(float3(input.texcoords, 1.0), " + stage.tcModName + ");");
                    }
                }
            }
            sb.AppendLine();
            sb.AppendLine("\treturn output;");
            sb.AppendLine("};");
            
            vertexShaderProgram = sb.ToString();
            sb.Remove(0, sb.ToString().Length);
            #endregion

            #region Generate PixelShaderProgram
            sb.AppendLine("float4 PixelShaderProgram(VertexShaderOutput input) : COLOR0");
            sb.AppendLine("{");
            sb.AppendLine("\tfloat4 finalColor = (float4)0;");
            foreach (Q3BSPMaterialStage stage in shader.stages)
            {
                if (stage.IsLightmapStage)
                {
                    sb.AppendLine("\tfloat4 lightmap = tex2D(" + stage.samplerName + ", input." + stage.stageName + ".uv);");
                    sb.AppendLine("\tfinalColor = finalColor*lightmap - finalColor * (1 - finalColor.a);");
                }
                else
                {
                    sb.AppendLine("\tfinalColor += tex2D(" + stage.samplerName + ", input." + stage.stageName + ".uv);");
                }
            }
            sb.AppendLine();
            sb.AppendLine("\treturn finalColor;");
            sb.AppendLine("}");

            pixelShaderProgram = sb.ToString();
            sb.Remove(0, sb.ToString().Length);
            #endregion

            #region Generate Technique
            sb.AppendLine("technique Technique1");
            sb.AppendLine("{");
            sb.AppendLine("\tpass Pass1");
            sb.AppendLine("\t{");
            sb.AppendLine("\t\tVertexShader = compile vs_1_1 VertexShaderProgram();");
            sb.AppendLine("\t\tPixelShader = compile ps_1_1 PixelShaderProgram();");
            sb.AppendLine("\t}");
            sb.AppendLine("}");

            technique = sb.ToString();
            sb.Remove(0, sb.ToString().Length);
            #endregion

            sw.Write(externs);
            sw.WriteLine();
            sw.Write(samplers);
            sw.WriteLine();
            sw.Write(stageInputStructs);
            sw.WriteLine();
            sw.Write(vertexShaderInput);
            sw.WriteLine();
            sw.Write(vertexShaderOutput);
            sw.WriteLine();
            sw.Write(vertexShaderProgram);
            sw.WriteLine();
            sw.Write(pixelShaderProgram);
            sw.WriteLine();
            sw.Write(technique);
            sw.WriteLine();
            sw.Flush();

            return externs + samplers + stageInputStructs + vertexShaderInput + vertexShaderOutput + vertexShaderProgram + pixelShaderProgram + technique;
        }

        void ProcessTcMod(string line, ref Q3BSPMaterialStage currentstage)
        {
            if (line.Length > 4 && line.Substring(0, 5) == "scale")
            {
                string[] split = line.Split(' ');
                currentstage.Scale = new Vector2(float.Parse(split[1]),float.Parse(split[2]));
            }

            else if (line.Length > 5 && line.Substring(0, 6) == "scroll")
            {
                string[] split = line.Split(' ');
                //currentstage.Scroll = new Vector2(float.Parse(split[1]), float.Parse(split[2]));
                currentstage.Scroll = new Vector2(1f, 1f);
            }

            else if (line.Length > 5 && line.Substring(0, 6) == "rotate")
            {
                string[] split = line.Split(' ');
                currentstage.Rotation = MathHelper.ToRadians(float.Parse(split[1]));
            }
        }

        private void ProcessSkyparms(string line, ref Q3BSPMaterial shader)
        {
            shader.IsSky = true;

            string[] split = line.Split(' ');
            if (split[0] != "-")
            {
                shader.FarBoxName = split[0];
            }
            if (split[1] != "-")
            {
                shader.SkyHeight = int.Parse(split[1]);
            }
            if (split[2] != "-")
            {
                shader.NearBoxName = split[2];
            }

            
        }

    }
}
